from .nodes import *
from .transitions import *
from .pilot import *
from .rrt import StartCollides, GoalCollides

class CheckHolding(StateNode):
    def start(self,event=None):
        super().start(event)
        if self.robot.holding is not None:
            self.post_success()
        else:
            self.post_failure()

class ClearHolding(StateNode):
    def start(self,event=None):
        super().start(event)
        held = self.robot.holding
        if held:
            held.held_by = None
            self.robot.holding = None

class PickUp(ActionNode,ObjectSpecNode):
    def __init__(self, object_spec=None):
        super().__init__()
        self.object_spec = object_spec

    def start(self,event=None):
        if isinstance(event, DataEvent):
            spec = event.data
        else:
            spec = self.object_spec
        self.object = self.get_object_from_spec(spec)
        super().start(event)
        if self.object is None:
            self.post_failure()

    class SetHolding(StateNode):
        def start(self,event=None):
            super().start(event)
            self.robot.holding = self.parent.object
            self.parent.object.held_by = self.robot

    class Pickup_PilotToObject(PilotToObject):
        def start(self,event=None):
            self.object = self.parent.object
            super().start(event)

    class Pickup_Forward(Forward):
        def start(self, event=None):
            if self.parent.object.is_visible:
                robot_radius = 30
                dx = self.parent.object.pose.x - self.robot.pose.x
                dy = self.parent.object.pose.y - self.robot.pose.y
                self.distance_mm = sqrt(dx**2 + dy**2) - robot_radius + 5
            else:
                self.distance_mm = 15
            super().start(event)

    def setup(self):
        #         # check StartCollides
        #         # drop object if carrying
        #         # path plan to object
        #         # check object is visible; fail if not
        #         # pick up
        #         # verify object held
        # 
        #         check_start: PilotCheckStart()
        #         check_start =F=> Print("Start collides!") =N=> Forward(-50) =C=> check_holding
        #         check_start =S=> check_holding
        # 
        #         check_holding: CheckHolding()
        #         check_holding =F=> plan
        #         check_holding =S=> Drop() =C=> plan
        # 
        #         plan: self.Pickup_PilotToObject()
        #         plan =C=> forward
        #         plan =PILOT(GoalCollides)=> ParentPilotEvent()
        # 
        #         forward: self.Pickup_Forward()
        #         forward =C=> self.SetHolding() =N=> ParentCompletes()
        #         forward =F=> ParentFails()
        
        # Code generated by genfsm on Wed Apr 23 00:28:21 2025:
        
        check_start = PilotCheckStart() .set_name("check_start") .set_parent(self)
        print1 = Print("Start collides!") .set_name("print1") .set_parent(self)
        forward1 = Forward(-50) .set_name("forward1") .set_parent(self)
        check_holding = CheckHolding() .set_name("check_holding") .set_parent(self)
        drop1 = Drop() .set_name("drop1") .set_parent(self)
        plan = self.Pickup_PilotToObject() .set_name("plan") .set_parent(self)
        parentpilotevent1 = ParentPilotEvent() .set_name("parentpilotevent1") .set_parent(self)
        forward = self.Pickup_Forward() .set_name("forward") .set_parent(self)
        setholding1 = self.SetHolding() .set_name("setholding1") .set_parent(self)
        parentcompletes1 = ParentCompletes() .set_name("parentcompletes1") .set_parent(self)
        parentfails1 = ParentFails() .set_name("parentfails1") .set_parent(self)
        
        failuretrans1 = FailureTrans() .set_name("failuretrans1")
        failuretrans1 .add_sources(check_start) .add_destinations(print1)
        
        nulltrans1 = NullTrans() .set_name("nulltrans1")
        nulltrans1 .add_sources(print1) .add_destinations(forward1)
        
        completiontrans1 = CompletionTrans() .set_name("completiontrans1")
        completiontrans1 .add_sources(forward1) .add_destinations(check_holding)
        
        successtrans1 = SuccessTrans() .set_name("successtrans1")
        successtrans1 .add_sources(check_start) .add_destinations(check_holding)
        
        failuretrans2 = FailureTrans() .set_name("failuretrans2")
        failuretrans2 .add_sources(check_holding) .add_destinations(plan)
        
        successtrans2 = SuccessTrans() .set_name("successtrans2")
        successtrans2 .add_sources(check_holding) .add_destinations(drop1)
        
        completiontrans2 = CompletionTrans() .set_name("completiontrans2")
        completiontrans2 .add_sources(drop1) .add_destinations(plan)
        
        completiontrans3 = CompletionTrans() .set_name("completiontrans3")
        completiontrans3 .add_sources(plan) .add_destinations(forward)
        
        pilottrans1 = PilotTrans(GoalCollides) .set_name("pilottrans1")
        pilottrans1 .add_sources(plan) .add_destinations(parentpilotevent1)
        
        completiontrans4 = CompletionTrans() .set_name("completiontrans4")
        completiontrans4 .add_sources(forward) .add_destinations(setholding1)
        
        nulltrans2 = NullTrans() .set_name("nulltrans2")
        nulltrans2 .add_sources(setholding1) .add_destinations(parentcompletes1)
        
        failuretrans3 = FailureTrans() .set_name("failuretrans3")
        failuretrans3 .add_sources(forward) .add_destinations(parentfails1)
        
        return self

################ Drop ################

class Drop(StateNode):
    def setup(self):
        #         check: CheckHolding()
        #         check =F=> ParentFails()
        #         check =S=> kick
        # 
        # 
        #         kick: SoftKick() =C=> ClearHolding() =N=> look
        # 
        #         # 1 second delay to give vision time to register the object
        #         look: StateNode() =T(1)=> done
        # 
        #         done: ParentCompletes()
        
        # Code generated by genfsm on Wed Apr 23 00:28:21 2025:
        
        check = CheckHolding() .set_name("check") .set_parent(self)
        parentfails2 = ParentFails() .set_name("parentfails2") .set_parent(self)
        kick = SoftKick() .set_name("kick") .set_parent(self)
        clearholding1 = ClearHolding() .set_name("clearholding1") .set_parent(self)
        look = StateNode() .set_name("look") .set_parent(self)
        done = ParentCompletes() .set_name("done") .set_parent(self)
        
        failuretrans4 = FailureTrans() .set_name("failuretrans4")
        failuretrans4 .add_sources(check) .add_destinations(parentfails2)
        
        successtrans3 = SuccessTrans() .set_name("successtrans3")
        successtrans3 .add_sources(check) .add_destinations(kick)
        
        completiontrans5 = CompletionTrans() .set_name("completiontrans5")
        completiontrans5 .add_sources(kick) .add_destinations(clearholding1)
        
        nulltrans3 = NullTrans() .set_name("nulltrans3")
        nulltrans3 .add_sources(clearholding1) .add_destinations(look)
        
        timertrans1 = TimerTrans(1) .set_name("timertrans1")
        timertrans1 .add_sources(look) .add_destinations(done)
        
        return self
